# -*- coding:UTF-8 -*-
import json
import os
from fastapi import APIRouter,Body,UploadFile,Form,File
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Union,List,Optional,Callable,Generator,Dict

from utils import BaseResponse,FILES_DATA_DIR,logger,retriveal_qdrant,FILE_AND_DOC_ID_STORGE
from data_loader.file_data_loader import FileLoader


file_app=APIRouter(tags=["文件管理接口"],prefix="/file")

def run_in_thread_pool(
        func: Callable,
        params: List[Dict] = [],
) -> Generator:
    '''
    在线程池中批量运行任务，并将运行结果以生成器的形式返回。
    请确保任务中的所有操作是线程安全的，任务函数请全部使用关键字参数。
    '''
    tasks = []
    with ThreadPoolExecutor() as pool:
        for kwargs in params:
            thread = pool.submit(func, **kwargs)
            tasks.append(thread)

        for obj in as_completed(tasks):
            yield obj.result()

# 多线程保存文件函数
def _save_files_in_thread(files: List[UploadFile],
                          override: bool):
    """
    通过多线程将上传的文件保存到对应知识库目录内。
    生成器返回保存结果：{"code":200, "msg": "xxx", "data": {"file_name": "xxx"}}
    """

    def save_file(file: UploadFile, override: bool=True) -> dict:
        '''
        保存单个文件。
        '''
        try:
            filename = file.filename
            file_path = os.path.join(FILES_DATA_DIR,filename)
            data = {"file_name": filename}
            file_content = file.file.read()  # 读取上传文件的内容
            if (os.path.isfile(file_path)
                    and not override
                    and os.path.getsize(file_path) == len(file_content)
            ):
                file_status = f"文件 {filename} 已存在。"
                logger.warn(file_status)
                return dict(code=404, msg=file_status, data=data)

            # 创建上传目录
            if not os.path.isdir(os.path.dirname(file_path)):
                os.makedirs(os.path.dirname(file_path))
            
            with open(file_path, "wb") as f:
                f.write(file_content)
            return dict(code=200, msg=f"成功上传文件 {filename}", data=data)
        except Exception as e:
            msg = f"{filename} 文件上传失败，报错信息为: {e}"
            logger.error(f'{e.__class__.__name__}: {msg}',
                         exc_info=e)
            return dict(code=500, msg=msg, data=data)

    params = [{"file": file, "override": override} for file in files]
    for result in run_in_thread_pool(save_file, params=params):
        yield result



def update_docs(
        file_names:List[str]
)->Dict:

    failed_files=[]

    for file_name in file_names:
        # 创建加载器
        try:
            file_loader=FileLoader(file_path=os.path.join(FILES_DATA_DIR,file_name))
            documnets=file_loader.load()
            ids=retriveal_qdrant.add_docments(documnets)
            FILE_AND_DOC_ID_STORGE[file_name]=ids
        except Exception as e:
            logger.error(f"文件 {file_name} 向量失败==>"+str(e))
            failed_files.append(file_name)
    return {"failed_files": failed_files}




@file_app.post("/upload_files",description="上传文档（可多个上传）")
async def upload_docs(
        files:List[UploadFile]=File(...,description="上传文件，支持多个文件"),
        override:bool=Form(True,description="是否覆盖已有的文件,默认为False"),
        is_vector_store:bool=Form(True,description="上传文件后是否进行向量化")
)->BaseResponse:
    
    failed_files=[] # 失败上传的文件
    success_files=set()
    for result in _save_files_in_thread(files,override):
        filename=result['data']["file_name"]
        if result["code"]!=200:
            failed_files.append(result['data'])
        if filename not in success_files:
            success_files.add(filename)

    # 对文件进行向量化
    if is_vector_store:
        ## 上传文档
        result=update_docs(
            file_names=success_files,
        )
        failed_files=failed_files+result["failed_files"]

    return BaseResponse(code=200, msg="文件上传与向量化完成", data={"failed_files": failed_files})

@file_app.get("/list",description="获取文件列表")
async def file_list():
    data=[]
    for root,_,files in os.walk(FILES_DATA_DIR):
        for file in files:
            data.append({
                "file_path":os.path.join(root,file),
                "file_name":file
            })
    return BaseResponse(data=data)

@file_app.post("/delete/{file_name}",description="获取文件列表")
async def delete_file(file_name:str):
    file_path=os.path.join(FILES_DATA_DIR,file_name)
    if not os.path.exists(file_path):
        return BaseResponse(code=500,msg=f"文件 {file_name} 不存在")
    if FILE_AND_DOC_ID_STORGE[file_name]:
        ids=FILE_AND_DOC_ID_STORGE[file_name]
        flag = retriveal_qdrant.delte_docment_by_ids(ids)
        if not flag:
                return BaseResponse(msg=f"{file_name} 向量库数据删除失败")
        FILE_AND_DOC_ID_STORGE.remove(file_name)
    os.remove(file_path)
    return BaseResponse(msg=f"文件 {file_name} 删除成功")